{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "45b3fe18-3361-4c79-9045-fa0dadc89d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9b36d415-c6bd-4912-a3f8-da2fba5d2f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = ConversationBufferMemory(return_messages=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "96cb92a1-4942-4677-a36c-d607bd650690",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': []}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "276127be-3191-4e11-ab44-ecfaef46ac02",
   "metadata": {},
   "outputs": [],
   "source": [
    "memory.save_context({\"input\": \"我的名字是林粒粒\"}, {\"output\": \"你好，林粒粒\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9ad55d0f-18e1-4aa5-86ef-38bf84971f7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [HumanMessage(content='我的名字是林粒粒'), AIMessage(content='你好，林粒粒')]}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f2a2ff6a-4bcd-4e5d-a369-a4e9037a5159",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [HumanMessage(content='我的名字是林粒粒'),\n",
       "  AIMessage(content='你好，林粒粒'),\n",
       "  HumanMessage(content='我是一名程序员'),\n",
       "  AIMessage(content='好的，我记住了')]}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.save_context({\"input\": \"我是一名程序员\"}, {\"output\": \"好的，我记住了\"})\n",
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f25d19c2-f01e-48f8-9ee5-71532282dd18",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"你是一个乐于助人的助手。\"),\n",
    "        MessagesPlaceholder(variable_name=\"history\"),\n",
    "        (\"human\", \"{user_input}\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fe4980f9-073c-48f2-a68c-3e2e7ed50a55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [HumanMessage(content='我的名字是林粒粒'),\n",
       "  AIMessage(content='你好，林粒粒'),\n",
       "  HumanMessage(content='我是一名程序员'),\n",
       "  AIMessage(content='好的，我记住了')]}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0d7801f4-510d-4036-b775-19e5645bf231",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatOpenAI(model=\"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "68c45a32-34c0-4c01-982a-b54172a4e640",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = prompt | model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1eb8738d-f226-4cf1-be48-7f1e41f91d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_input = \"你知道我的名字吗？\"\n",
    "history = memory.load_memory_variables({})[\"history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1606e9c7-745d-4401-a608-34995c9a0fa3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='是的，你告诉我了，你的名字是林粒粒。')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = chain.invoke({\n",
    "    \"user_input\": user_input, \n",
    "    'history': history\n",
    "})\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "aa2456aa-b520-493d-9257-8f32cc27aa81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [HumanMessage(content='我的名字是林粒粒'),\n",
       "  AIMessage(content='你好，林粒粒'),\n",
       "  HumanMessage(content='我是一名程序员'),\n",
       "  AIMessage(content='好的，我记住了'),\n",
       "  HumanMessage(content='你知道我的名字吗？'),\n",
       "  AIMessage(content='是的，你告诉我了，你的名字是林粒粒。')]}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.save_context({\"input\": user_input}, {\"output\": result.content})\n",
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "967bee41-0f62-44ba-aa3a-71eee83cf942",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='你上一个问题问我：\"你知道我的名字吗？\"')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_input = \"根据对话历史告诉我，我上一个问题问你的是什么？请重复一遍\"\n",
    "history = memory.load_memory_variables({})[\"history\"]\n",
    "\n",
    "result = chain.invoke({\n",
    "    \"user_input\": user_input, \n",
    "    'history': history\n",
    "})\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "840f2919-0eda-408e-86b2-ecbc24209e75",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
